import pickle
import tensorflow as tf
import numpy as np
import os
from sklearn.model_selection import train_test_split

BOTTLENECK_TENSOR_NAME = 'pool_3/_reshape:0'
BOTTLENECK_VEC_SIZE = 2048
JPEG_DATA_TENSOR_NAME = 'DecodeJpeg/contents:0'
MODEL_PATH = '../../../../large_data/model/inceptionV3/tensorflow_inception_graph.pb'
IMG_ROOT_DIR = '../../../../large_data/CV2/_many_files/flower_photos_liuqilong/'
PERCENT_TRAIN = 80
PERCENT_TEST = 10
PERCENT_VAL = 10


def sep(label='', cnt=32):
    print('-' * cnt, label, '-' * cnt, sep='')


########################################################################################################################
sep('Prepare')
tf.random.set_random_seed(777)
np.random.seed(777)

VER = 'v7.0'
FILE_NAME = os.path.basename(__file__)
SAVE_DIR = os.path.join('_save', FILE_NAME, VER)
SAVE_NAME = 'bottleneck.pickle'
os.makedirs(SAVE_DIR, exist_ok=True)
SAVE_PATH = os.path.join(SAVE_DIR, SAVE_NAME)

########################################################################################################################
sep('Load model')
# 读取模型
with open(MODEL_PATH, 'rb') as f:
    model_data = f.read()
graph_def = tf.GraphDef()
graph_def.ParseFromString(model_data)
# 从模型种获取张量
input_placeholder, output_tensor = tf.import_graph_def(graph_def,
                                                       return_elements=[
                                                           JPEG_DATA_TENSOR_NAME,
                                                           BOTTLENECK_TENSOR_NAME,
                                                       ])

########################################################################################################################
sep('Load pictures')
sess = tf.Session()  # 为迁移学习获取特征开启会话
label_idx = 0
ext_set = set(['.jpg', '.jpeg'])
x_data = []
y_data = []
relative_path_data = []
idx2name = {}
name2idx = {}
cnt = 0
print('正在加载……')
for dir_name in os.listdir(IMG_ROOT_DIR):
    dir_path = os.path.join(IMG_ROOT_DIR, dir_name)

    # 如果不是文件夹
    if not os.path.isdir(dir_path):
        continue

    # 列举文件夹里的条目
    file_names = os.listdir(dir_path)
    # 如果是空文件夹
    if len(file_names) == 0:
        continue

    idx2name[label_idx] = dir_name
    name2idx[dir_name] = label_idx

    # 检查每个条目
    for file_name in file_names:
        file_path = os.path.join(dir_path, file_name)
        # 是目录
        if os.path.isdir(file_path):
            continue
        # 扩展名是否是图片
        ext = os.path.splitext(file_name)[1].lower()
        if not ext in ext_set:
            continue
        # 读取图片
        with open(file_path, 'rb') as f:
            img_data = f.read()

        # 获取特征
        vec = sess.run(output_tensor, feed_dict={input_placeholder: img_data})
        # 1行矩阵转向量
        vec = np.squeeze(vec, axis=0)

        # 保存
        x_data.append(vec)
        y_data.append(label_idx)
        relative_path_data.append(dir_name + '/' + file_name)

        # 计数
        cnt += 1
        if cnt % 25 == 0:
            print(f'已经加载{cnt}张图片。')

    label_idx += 1

print(f'加载完成！共加载{cnt}张图片。')

########################################################################################################################
sep('Split data')
# 注意：这里用的库函数带乱序。如果库函数没有乱序，一定要主动乱序！
x_train, x_test_val, y_train, y_test_val, path_train, path_test_val \
    = train_test_split(x_data, y_data, relative_path_data,
                       train_size=PERCENT_TRAIN/100,
                       shuffle=True,
                       random_state=777)
x_test, x_val, y_test, y_val, path_test, path_val \
    = train_test_split(x_test_val ,y_test_val, path_test_val,
                       train_size=PERCENT_TEST/(PERCENT_TEST + PERCENT_VAL),
                       shuffle=True,
                       random_state=777)

print('x_train', np.shape(x_train))
print('y_train', np.shape(y_train))
print('path_train', np.shape(path_train))
print('x_test', np.shape(x_test))
print('y_test', np.shape(y_test))
print('path_test', np.shape(path_test))
print('x_val', np.shape(x_val))
print('y_val', np.shape(y_val))
print('path_val', np.shape(path_val))

########################################################################################################################
sep('Pickle data')
data = {
    'idx2name': idx2name,
    'name2idx': name2idx,
    'x_train': x_train,
    'y_train': y_train,
    'path_train': path_train,
    'x_test': x_test,
    'y_test': y_test,
    'path_test': path_test,
    'x_val': x_val,
    'y_val': y_val,
    'path_val': path_val,
}
with open(SAVE_PATH, 'wb') as f:
    pickle.dump(data, f)
print(f'Data saved by pickle to {SAVE_PATH}')

########################################################################################################################
sep('End the job')
sess.close()  # 最后关闭会话

########################################################################################################################
sep('Test saved data')
with open(SAVE_PATH, 'rb') as f:
    data = pickle.load(f)
for key in data.keys():
    print(f'{key}: ', end='')
    item = data[key]
    if type(item) == dict:
        print(item)
    else:
        print(np.shape(item))

########################################################################################################################
sep('Over')
